import os

from bson import ObjectId
from loguru import logger
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import UpdateOne

from config import settings
from utils.dateutil import date_now
from utils.dbutil import convert2sort, convert2projection
from .basic import mdb


async def connect_to_mongo():
    mdb.clients = dict()
    for name, mongo_conf in settings.mongo.items():
        logger.info(f"init mongodb client {name}")
        host = os.getenv("MONGO_URL") or mongo_conf['host']
        mdb.clients[name.lower()] = AsyncIOMotorClient("mongodb://admin:kejia@42.192.110.104:27017/?authMechanism=DEFAULT")


async def close_mongo_connection():
    for name, client in mdb.clients.items():
        logger.info(f"close mongodb client {name}")
        client.close()


async def find_one_doc(col_const, query, **kwargs):
    col = mdb.get_collection(col_const)
    result = await col.find_one(query, **kwargs)
    return result


async def find_many_docs(col_const, query, **kwargs):
    col = mdb.get_collection(col_const)
    results = [i async for i in col.find(query, **kwargs)]
    return results


async def update_one_doc(col_const, query, update):
    col = mdb.get_collection(col_const)
    result = await col.update_one(query, update)
    return result.modified_count


async def update_docs(col_const, query, update):
    col = mdb.get_collection(col_const)
    result = await col.update_many(query, update)
    return result.modified_count


async def aggregate_docs(col_const, pipeline):
    col = mdb.get_collection(col_const)
    result = list()
    async for doc in col.aggregate(pipeline):
        result.append(doc)
    return result


async def bulk_write(col_const, opts):
    col = mdb.get_collection(col_const)
    result = await col.bulk_write(opts)
    return result


async def insert_one_model(col_const, model):
    doc = model.mongo(exclude_unset=True)
    col = mdb.get_collection(col_const)
    result = await col.insert_one(doc)
    return result.inserted_id


async def insert_models(col_const, models):
    docs = [m.mongo(exclude_unset=True) for m in models]
    col = mdb.get_collection(col_const)
    result = await col.insert_many(docs)
    return result.inserted_ids


async def find_one_model_by_id(col_const, oid: str, model, fields=None):
    col = mdb.get_collection(col_const)
    query = {'_id': ObjectId(oid)}
    if fields:
        projection = convert2projection(fields)
        doc = await col.find_one(query, projection=projection)
    else:
        doc = await col.find_one(query)
    return model.from_mongo(doc)


async def find_one_model(col_const, query, model, fields=None, **kwargs):
    col = mdb.get_collection(col_const)
    sort = kwargs.get('sort', [])
    sort = convert2sort(sort)

    doc = await col.find_one(query, projection=fields, sort=sort)
    return model.from_mongo(doc)


async def find_many_models(col_const, model, query, sort, offset, size, fields=None):
    col = mdb.get_collection(col_const)
    projection = convert2projection(fields)
    sort = convert2sort(sort)
    cursor = col.find(query, projection=projection, sort=sort, skip=offset, limit=size)

    results = list()
    async for doc in cursor:
        results.append(model.from_mongo(doc))
    return results


async def find_mono_docs(col_const, query, sort, offset, size, fields=None):
    col = mdb.get_collection(col_const)
    projection = convert2projection(fields)
    sort = convert2sort(sort)
    cursor = col.find(query, projection=projection, sort=sort, skip=offset, limit=size)

    results = list()
    async for doc in cursor:
        results.append(doc)
    return results




async def aggregate_models(col_const, pipeline, model):
    col = mdb.get_collection(col_const)
    result = list()
    async for doc in col.aggregate(pipeline):
        result.append(model.from_mongo(doc))
    return result


async def count_documents(col_const, query):
    col = mdb.get_collection(col_const)
    num = await col.count_documents(query)
    return num


async def update_one_model_by_id(col_const, oid: str, updated_model):
    col = mdb.get_collection(col_const)
    query = {'_id': ObjectId(oid)}
    update = updated_model.mongo(exclude_unset=True)
    if not update:
        return 0
    update = {"$set": update}
    result = await col.update_one(query, update)
    return result.modified_count


async def update_one_model(col_const, query, update_model):
    col = mdb.get_collection(col_const)
    update = update_model.mongo(exclude_unset=True)
    if not update:
        return 0
    update = {"$set": update}
    result = await col.update_one(query, update)
    return result.modified_count


def warp_update(model, action):
    # todo  edit, all field_set overwrite[warning]
    _exclude_none = False
    if action == 'batches':
        # concern 有值的情况.
        _exclude_none = True
    doc = model.mongo(exclude_unset=True, exclude_none=_exclude_none, exclude={"num_of_activities", "belongs_inst"})
    oid = doc.pop("_id")
    return UpdateOne({'_id': oid}, {'$set': doc})


async def bulk_update_models_through_id(col_const, update_models, action="edit"):
    updated_docs = [warp_update(m, action) for m in update_models]
    result = await bulk_write(col_const, updated_docs)
    return result.modified_count


async def bulk_update_models_unset_fields_though_id(col_const, update_models, model_type, action=None):
    """
    model_type: 最终需要上传的model，可以和update_models的model不一致，但是字段必须覆盖
    # update_models 更新的数据
    # model_type Model Class
    """
    query = {"_id": {"$in": [ObjectId(update_model.id) for update_model in update_models]}}
    results = await find_many_models(col_const, model_type, query=query, offset=0, size=0, sort=None)
    exists_models = {str(res.id): res for res in results}
    update_items = []
    for update_model in update_models:
        exists_model = exists_models[str(update_model.id)]
        exclude_fields = [key for key, value in exists_model.dict(exclude_unset=True, exclude={"id"}).items() if
                          value is not None]
        update_dict = update_model.dict(exclude=set(exclude_fields), exclude_unset=True)
        if not update_dict:
            continue

        update_model = model_type(**update_dict)
        update_items.append(update_model)
    if not update_items:
        return 0
    res = await bulk_update_models_through_id(col_const, update_items, action)
    return res


async def delete_one_model_by_id(col_const, oid):
    query = {'_id': ObjectId(oid)}
    return await delete_one_model(col_const, query)


async def delete_one_model(col_const, query):
    col = mdb.get_collection(col_const)
    result = await col.delete_one(query)
    return result.deleted_count


async def delete_models(col_const, query):
    col = mdb.get_collection(col_const)
    result = await col.delete_many(query)
    return result.deleted_count


async def find_one_doc_and_update(col_const, query, update, **kwargs):
    col = mdb.get_collection(col_const)
    result = await col.find_one_and_update(query, update, **kwargs)
    return result


async def fake_delete_models(col_const, query):
    update = {"$set": {"is_deleted": True, "deleted_time": date_now()}}
    return await update_docs(col_const, query, update)


async def fake_delete_models_by_ids(col_const, ids):
    query = {"_id": {"$in": [ObjectId(i) for i in ids]}}
    return await fake_delete_models(col_const, query)


async def check_models_exists_by_ids(col_const, ids, with_deleted=False):
    query = {
        "_id": {"$in": [ObjectId(i) for i in ids]},
        "is_deleted": True if with_deleted else {"$ne": True}
    }
    docs = await find_many_docs(col_const, query, projection=['_id'])
    exists = {str(d['_id']) for d in docs}
    missing = set(ids) - exists
    return list(missing)


async def lookup_oid_reference(models, from_field, to_field, refer_col, refer_model_type, refer_fields):
    refer_ids = []
    for m in models:
        refer_value = getattr(m, from_field)
        if refer_value is None:
            continue
        if isinstance(refer_value, list):
            refer_ids.extend(refer_value)
        else:
            refer_ids.append(refer_value)
    if not refer_ids:
        return
    query = {"_id": {"$in": refer_ids}}
    refer_models = await find_many_models(refer_col, refer_model_type, query, None, 0, 0, refer_fields)
    id2model = {m.id: m for m in refer_models}
    missing = {i for i in refer_ids if i not in id2model}
    logger.warning(f"lookup_oid_reference:{refer_model_type.__name__} missing {', '.join((str(i) for i in missing))}")
    for m in models:
        refer_value = getattr(m, from_field)
        if refer_value is None:
            continue
        if isinstance(refer_value, list):
            setattr(m, to_field, [id2model.get(i) for i in id2model])
        else:
            setattr(m, to_field, id2model.get(refer_value))


async def lookup_oid_reference_multi_fields_same_type(models, fields_mapping: dict, refer_col, refer_model_type,
                                                      refer_fields):
    """
    多个字段，引用同一种实体
    !!! TODO 需要和 lookup_oid_reference 尽量合并，是否可以写一个 LookupTool 的工具类进行方法的封装？
    :fields_mapping {from_file: to_file}
    :return:
    """
    refer_ids = []
    for m in models:
        for from_field in fields_mapping.keys():
            refer_value = getattr(m, from_field)
            if refer_value is None:
                continue
            if isinstance(refer_value, list):
                refer_ids.extend(refer_value)
            else:
                refer_ids.append(refer_value)
    if not refer_ids:
        return
    query = {"_id": {"$in": refer_ids}}
    refer_models = await find_many_models(refer_col, refer_model_type, query, None, 0, 0, refer_fields)
    id2model = {m.id: m for m in refer_models}
    missing = {i for i in refer_ids if i not in id2model}
    logger.warning(f"lookup_oid_reference:{refer_model_type.__name__} missing {', '.join((str(i) for i in missing))}")
    for m in models:
        for from_field, to_field in fields_mapping.items():
            refer_value = getattr(m, from_field)
            if refer_value is None:
                continue
            if isinstance(refer_value, list):
                setattr(m, to_field, [id2model.get(i) for i in refer_value])
            else:
                setattr(m, to_field, id2model.get(refer_value))
